import json
import paramiko
import logging
import threading
import time
import uuid
import select
import socket
import asyncio
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from django.contrib.auth.models import User
from django.utils import timezone
from .models import Host, ConnectionRecord, SSHSession

logger = logging.getLogger(__name__)


class NewSSHConsumer(AsyncWebsocketConsumer):
    """新的SSH WebSocket Consumer - 更稳定的实现"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ssh_client = None
        self.ssh_channel = None
        self.host_id = None
        self.host = None
        self.user = None
        self.connection_record = None
        self.session_record = None
        # 移除审计服务以简化
        self.read_thread = None
        self.is_connected = False
        self.terminal_size = {'cols': 80, 'rows': 24}

    async def connect(self):
        """WebSocket连接建立"""
        try:
            # 获取用户信息
            self.user = self.scope.get("user")
            # logger.info(f"WebSocket连接尝试，用户: {self.user}, 认证状态: {getattr(self.user, 'is_authenticated', False) if self.user else False}")  # 注释掉避免终端输出干扰

            if not self.user or not getattr(self.user, 'is_authenticated', False):
                logger.warning("WebSocket连接被拒绝：用户未认证")
                await self.close(code=4001)
                return

            # 获取主机ID
            self.host_id = self.scope['url_route']['kwargs']['host_id']
            # logger.info(f"WebSocket连接主机ID: {self.host_id}")  # 注释掉避免终端输出干扰

            # 验证访问令牌
            query_string = self.scope.get('query_string', b'').decode('utf-8')
            token = None
            if query_string:
                from urllib.parse import parse_qs
                query_params = parse_qs(query_string)
                token = query_params.get('token', [None])[0]

            if not token:
                logger.warning("WebSocket连接被拒绝：缺少访问令牌")
                await self.close(code=4002)
                return

            # 验证令牌有效性
            if not await self._validate_access_token(token):
                logger.warning("WebSocket连接被拒绝：访问令牌无效")
                await self.close(code=4003)
                return

            # 接受WebSocket连接
            await self.accept()
            # logger.info("WebSocket连接已接受")  # 注释掉避免终端输出干扰

            # 发送连接状态
            await self.send(text_data=json.dumps({
                'type': 'status',
                'status': 'websocket_connected',
                'message': 'WebSocket连接已建立'
            }))

        except Exception as e:
            logger.error(f"WebSocket连接失败: {e}")
            await self.close(code=4000)

    async def disconnect(self, close_code):
        """WebSocket连接断开"""
        await self.cleanup_ssh_connection()

    async def receive(self, text_data):
        """接收WebSocket消息"""
        try:
            data = json.loads(text_data)
            message_type = data.get('type')

            if message_type == 'auth':
                # 处理认证和连接
                ssh_port = data.get('ssh_port', 22)
                ssh_username = data.get('ssh_username')
                ssh_password = data.get('ssh_password')
                terminal_size = data.get('terminal_size', {'cols': 80, 'rows': 24})
                environment = data.get('environment', {})

                self.terminal_size = terminal_size
                self.environment = environment
                self.ssh_port = ssh_port
                self.ssh_username = ssh_username
                self.ssh_password = ssh_password

                await self.establish_ssh_connection()

            elif message_type == 'connect':
                # 建立SSH连接（兼容旧版本）
                await self.establish_ssh_connection()

            elif message_type == 'input' and self.ssh_channel:
                # 发送用户输入
                input_data = data.get('data', '')
                await self.send_input_to_ssh(input_data)

            elif message_type == 'resize':
                # 调整终端大小
                cols = data.get('cols', 80)
                rows = data.get('rows', 24)
                await self.resize_terminal(cols, rows)

            elif message_type == 'disconnect':
                # 断开SSH连接
                await self.cleanup_ssh_connection()

        except json.JSONDecodeError:
            logger.error("接收到无效的JSON数据")
        except Exception as e:
            logger.error(f"处理WebSocket消息失败: {e}")

    async def establish_ssh_connection(self):
        """建立SSH连接"""
        try:
            # 获取主机信息
            self.host = await self.get_host(self.host_id)
            if not self.host:
                await self.send_error("主机不存在")
                return

            # 检查权限
            if not await self.check_permission():
                await self.send_error("无权访问此主机")
                return

            # 创建连接记录
            self.connection_record = await self.create_connection_record()

            # 在线程池中建立SSH连接
            import asyncio
            await asyncio.get_event_loop().run_in_executor(
                None, self._create_ssh_connection
            )

            if self.is_connected:
                # 创建会话记录
                self.session_record = await self.create_session_record()

                # 启动输出读取线程
                self.read_thread = threading.Thread(target=self._read_ssh_output)
                self.read_thread.daemon = True
                self.read_thread.start()

                # 启动输出发送任务
                asyncio.create_task(self._process_output_queue())

                # 立即读取一次输出以获取初始提示符
                await asyncio.sleep(0.2)  # 等待一下让SSH连接稳定
                await self._read_initial_output()

                await self.send(text_data=json.dumps({
                    'type': 'connected',
                    'status': 'success',
                    'message': f'成功连接到 {self.host.hostname}',
                    'connection_id': self.connection_record.id if self.connection_record else None
                }))
            else:
                await self.send_error("SSH连接失败")

        except Exception as e:
            logger.error(f"建立SSH连接失败: {e}")
            await self.send_error(f"连接失败: {str(e)}")

    def _create_ssh_connection(self):
        """在线程中创建SSH连接"""
        try:
            self.ssh_client = paramiko.SSHClient()
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            # 处理localhost特殊情况
            hostname = self.host.public_ip
            if hostname in ['localhost', '127.0.0.1', '::1']:
                hostname = '127.0.0.1'

            logger.info(f"尝试SSH连接: {hostname}:{self.host.port}")

            # 根据认证方式连接
            try:
                if hasattr(self, 'ssh_username') and hasattr(self, 'ssh_password'):
                    # 使用前端传来的认证信息
                    logger.info(f"使用前端认证信息连接: 用户={self.ssh_username}, 端口={self.ssh_port}")
                    self.ssh_client.connect(
                        hostname=hostname,
                        port=self.ssh_port,
                        username=self.ssh_username,
                        password=self.ssh_password,
                        timeout=15,  # 增加超时时间
                        allow_agent=False,
                        look_for_keys=False,
                        banner_timeout=30  # 增加banner超时时间
                    )
                elif self.host.auth_type == 'password':
                    logger.info(f"使用密码认证连接: 用户={self.host.username}, 端口={self.host.port}")
                    self.ssh_client.connect(
                        hostname=hostname,
                        port=self.host.port,
                        username=self.host.username,
                        password=self.host.password,
                        timeout=15,
                        allow_agent=False,
                        look_for_keys=False,
                        banner_timeout=30
                    )
                elif self.host.auth_type == 'key':
                    logger.info(f"使用密钥认证连接: 用户={self.host.username}, 端口={self.host.port}")
                    from io import StringIO
                    private_key = paramiko.RSAKey.from_private_key(
                        StringIO(self.host.private_key),
                        password=self.host.passphrase if self.host.passphrase else None
                    )
                    self.ssh_client.connect(
                        hostname=hostname,
                        port=self.host.port,
                        username=self.host.username,
                        pkey=private_key,
                        timeout=15,
                        allow_agent=False,
                        look_for_keys=False,
                        banner_timeout=30
                    )
                else:
                    raise Exception("不支持的认证方式")

            except paramiko.ssh_exception.SSHException as e:
                error_msg = str(e)
                if "Error reading SSH protocol banner" in error_msg:
                    raise Exception(f"SSH连接被拒绝 - 可能原因：1) SSH服务未运行 2) 端口{self.host.port}被防火墙阻止 3) 主机{hostname}不可达。详细错误: {error_msg}")
                elif "Connection reset by peer" in error_msg:
                    raise Exception(f"连接被目标主机重置 - 可能原因：1) SSH服务配置拒绝连接 2) 防火墙规则 3) 网络问题。请检查主机{hostname}:{self.host.port}的SSH服务状态")
                else:
                    raise Exception(f"SSH协议错误: {error_msg}")
            except socket.timeout:
                raise Exception(f"连接超时 - 主机{hostname}:{self.host.port}无响应，请检查网络连接和防火墙设置")
            except socket.error as e:
                if e.errno == 111:  # Connection refused
                    raise Exception(f"连接被拒绝 - 主机{hostname}端口{self.host.port}未开放或SSH服务未运行")
                elif e.errno == 113:  # No route to host
                    raise Exception(f"无法到达主机{hostname} - 请检查网络路由和主机状态")
                else:
                    raise Exception(f"网络连接错误: {e}")
            except Exception as e:
                if "Authentication failed" in str(e):
                    raise Exception(f"认证失败 - 请检查用户名和密码是否正确")
                else:
                    raise e

            # 创建交互式shell - 使用前端传来的环境变量
            environment = getattr(self, 'environment', {})
            if not environment:
                # 如果没有前端环境变量，使用默认值
                environment = {
                    'TERM': 'xterm-256color',
                    'LANG': 'en_US.UTF-8',
                    'LC_ALL': 'en_US.UTF-8',
                    'COLUMNS': str(self.terminal_size['cols']),
                    'LINES': str(self.terminal_size['rows']),
                    'LS_COLORS': 'rs=0:di=01;34:ln=01;36:mh=00:pi=40;33:so=01;35:do=01;35:bd=40;33;01:cd=40;33;01:or=40;31;01:mi=00:su=37;41:sg=30;43:ca=30;41:tw=30;42:ow=34;42:st=37;44:ex=01;32',
                    'COLORTERM': 'truecolor',
                    'CLICOLOR': '1',
                    'CLICOLOR_FORCE': '1',
                    'LSCOLORS': 'ExFxCxDxBxegedabagacad'
                }

            self.ssh_channel = self.ssh_client.invoke_shell(
                term='xterm-256color',
                width=self.terminal_size['cols'],
                height=self.terminal_size['rows'],
                environment=environment
            )

            # 设置终端为原生模式
            self.ssh_channel.settimeout(0.1)  # 设置超时

            # 等待一下让连接稳定
            time.sleep(0.2)

            # 发送回车以显示提示符
            self.ssh_channel.send(b'\r\n')

            # 再等待一下让提示符显示
            time.sleep(0.1)

            self.is_connected = True
            logger.info(f"SSH连接成功: {self.host.hostname}")

        except Exception as e:
            logger.error(f"SSH连接失败: {e}")
            self.is_connected = False
            raise e




    def _read_ssh_output(self):
        """读取SSH输出的线程函数 - 优化版本"""
        import select
        import queue

        # 创建输出队列
        if not hasattr(self, 'output_queue'):
            self.output_queue = queue.Queue(maxsize=5000)  # 增加队列大小以匹配大缓冲区

        buffer = b''  # 字节缓冲区

        while self.is_connected and self.ssh_channel:
            try:
                # 使用select检查是否有数据可读
                ready, _, _ = select.select([self.ssh_channel], [], [], 0.1)

                if ready:
                    try:
                        # 使用65535字节的大缓冲区
                        data = self.ssh_channel.recv(65535)
                        if data:
                            buffer += data

                            # 按行分割输出，避免大块数据阻塞
                            while b'\n' in buffer:
                                line, buffer = buffer.split(b'\n', 1)
                                try:
                                    output_line = line.decode('utf-8', errors='replace') + '\n'
                                    self._send_to_websocket(output_line)
                                    # 添加小延迟避免前端渲染阻塞
                                    time.sleep(0.001)
                                except queue.Full:
                                    logger.warning("输出队列已满，丢弃部分数据")
                                    break

                            # 处理剩余的不完整行（如果缓冲区太大则强制输出）
                            if len(buffer) > 65535:  # 65KB限制
                                try:
                                    output_chunk = buffer.decode('utf-8', errors='replace')
                                    self._send_to_websocket(output_chunk)
                                    buffer = b''
                                except queue.Full:
                                    logger.warning("输出队列已满，清空缓冲区")
                                    buffer = b''
                        else:
                            # 连接关闭
                            break
                    except Exception as e:
                        logger.error(f"接收SSH数据失败: {e}")
                        break

            except Exception as e:
                logger.error(f"SSH输出读取失败: {e}")
                break

        # 处理剩余缓冲区数据
        if buffer:
            try:
                final_output = buffer.decode('utf-8', errors='replace')
                self._send_to_websocket(final_output)
            except:
                pass

        logger.info("SSH输出读取线程结束")

    def _send_to_websocket(self, data):
        """发送数据到WebSocket - 简化版本"""
        try:
            # 将数据放入队列，由主线程处理
            if not hasattr(self, 'output_queue'):
                import queue
                self.output_queue = queue.Queue()

            self.output_queue.put(data)

        except Exception as e:
            logger.error(f"发送到WebSocket失败: {e}")

    async def _process_output_queue(self):
        """处理输出队列的异步任务"""
        import queue

        while self.is_connected:
            try:
                if hasattr(self, 'output_queue'):
                    try:
                        # 非阻塞获取队列数据
                        output = self.output_queue.get_nowait()
                        await self.send_output_to_websocket(output)
                        self.output_queue.task_done()
                    except queue.Empty:
                        # 队列为空，短暂等待
                        await asyncio.sleep(0.01)
                else:
                    await asyncio.sleep(0.1)
            except Exception as e:
                logger.error(f"处理输出队列失败: {e}")
                await asyncio.sleep(0.1)

    async def send_output_to_websocket(self, output):
        """发送输出到WebSocket"""
        try:
            await self.send(text_data=json.dumps({
                'type': 'output',
                'data': output
            }))
        except Exception as e:
            logger.error(f"发送输出到WebSocket失败: {e}")

    async def send_input_to_ssh(self, input_data):
        """发送输入到SSH"""
        try:
            if self.ssh_channel and self.is_connected:
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_channel.send, input_data.encode('utf-8')
                )

                # 审计日志已简化

        except Exception as e:
            logger.error(f"发送输入到SSH失败: {e}")

    async def resize_terminal(self, cols, rows):
        """调整终端大小"""
        try:
            self.terminal_size = {'cols': cols, 'rows': rows}
            if self.ssh_channel and self.is_connected:
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_channel.resize_pty, cols, rows
                )

                # 更新环境变量以确保vim等应用获得正确的终端尺寸
                await asyncio.get_event_loop().run_in_executor(
                    None, self._update_terminal_env, cols, rows
                )

                logger.info(f"终端尺寸已调整为: {cols}x{rows}")
        except Exception as e:
            logger.error(f"调整终端大小失败: {e}")

    def _update_terminal_env(self, cols, rows):
        """更新终端环境变量 - 静默版本，不发送echo命令"""
        try:
            if self.ssh_channel and self.is_connected:
                # 只发送stty命令更新终端尺寸，不设置环境变量（避免echo）
                stty_cmd = f"stty cols {cols} rows {rows} >/dev/null 2>&1\n"
                self.ssh_channel.send(stty_cmd.encode('utf-8'))

        except Exception as e:
            logger.error(f"更新终端环境变量失败: {e}")

    async def _validate_access_token(self, token):
        """验证访问令牌"""
        try:
            from django.core.cache import cache

            cache_key = f"ssh_access_token:{token}"
            token_data = cache.get(cache_key)

            if not token_data:
                logger.warning(f"访问令牌无效或已过期: {token}")
                return False

            # 验证用户和主机匹配
            if (token_data['user_id'] != self.user.id or
                token_data['host_id'] != int(self.host_id)):
                logger.warning(f"访问令牌与用户或主机不匹配: user_id={token_data['user_id']}, host_id={token_data['host_id']}")
                return False

            logger.info(f"访问令牌验证成功: user_id={self.user.id}, host_id={self.host_id}")
            return True

        except Exception as e:
            logger.error(f"验证访问令牌失败: {e}")
            return False

    async def cleanup_ssh_connection(self):
        """清理SSH连接"""
        try:
            self.is_connected = False

            if self.ssh_channel:
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_channel.close
                )
                self.ssh_channel = None

            if self.ssh_client:
                await asyncio.get_event_loop().run_in_executor(
                    None, self.ssh_client.close
                )
                self.ssh_client = None

            # 更新连接记录
            if self.connection_record:
                await self.update_connection_record()

            # 更新会话记录
            if self.session_record:
                await self.update_session_record()

        except Exception as e:
            logger.error(f"清理SSH连接失败: {e}")

    async def send_error(self, message):
        """发送错误消息"""
        await self.send(text_data=json.dumps({
            'type': 'error',
            'status': 'error',
            'message': message
        }))

    # 数据库操作方法
    @database_sync_to_async
    def get_host(self, host_id):
        try:
            return Host.objects.get(id=host_id)
        except Host.DoesNotExist:
            return None

    @database_sync_to_async
    def check_permission(self):
        # 简单权限检查，可以根据需要扩展
        return self.user.is_authenticated

    @database_sync_to_async
    def create_connection_record(self):
        return ConnectionRecord.objects.create(
            user=self.user,
            host=self.host,
            is_active=True,
            connected_at=timezone.now()
        )

    @database_sync_to_async
    def create_session_record(self):
        return SSHSession.objects.create(
            user=self.user,
            host=self.host,
            session_id=str(uuid.uuid4()),
            status='active'
        )

    @database_sync_to_async
    def update_connection_record(self):
        if self.connection_record:
            self.connection_record.is_active = False
            self.connection_record.disconnected_at = timezone.now()
            self.connection_record.save()

    @database_sync_to_async
    def update_session_record(self):
        if self.session_record:
            self.session_record.status = 'closed'
            self.session_record.end_time = timezone.now()
            self.session_record.save()

    @database_sync_to_async
    def record_input(self, input_data):
        # 记录输入审计日志
        if self.audit_service and self.session_record:
            self.audit_service.record_command(
                self.session_record.session_id,
                input_data,
                'input'
            )

    @database_sync_to_async
    def record_output(self, output_data):
        # 记录输出审计日志
        if self.audit_service and self.session_record:
            self.audit_service.record_command(
                self.session_record.session_id,
                output_data,
                'output'
            )
